package cas.ihep.spark.test;

import cas.ihep.fs.UnifiedBlock;
import cas.ihep.fs.UnifiedFile;
import cas.ihep.fs.UnifiedFileSystem;
import cas.ihep.util.MemoryBuffer;
import com.google.common.io.Closer;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;

import java.io.*;
import java.net.URI;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;

public class WriteRandom {
    private static final int _1G=1024*1024*1024;

    private static long test1(int integer)throws Exception{
        Random rand=new Random(integer);
        long startTime;
        try(Closer closer=Closer.create()){
            UnifiedFileSystem fs=closer.register(UnifiedFileSystem.get(UnifiedFileSystem.newConfiguration(System.getenv("HADOOP_HOME"))));
            DataOutputStream fos=closer.register(new DataOutputStream(new BufferedOutputStream(fs.createStream(new URI("alluxio://hadoop06.ihep.ac.cn:19998/test/"+Integer.toHexString(integer))))));
            startTime=System.nanoTime();
            for(int i=0;i<_1G/4;i++){
                fos.writeInt(rand.nextInt());
            }
            startTime=System.nanoTime()-startTime;
        }
        return startTime;
    }

    private static long test2(int integer)throws Exception{
        Random rand=new Random(integer);
        long startTime;
        try(Closer closer=Closer.create()){
            UnifiedFileSystem fs=closer.register(UnifiedFileSystem.get(UnifiedFileSystem.newConfiguration(System.getenv("HADOOP_HOME"))));
            UnifiedFile file=closer.register(fs.create(new URI("alluxio://hadoop06.ihep.ac.cn:19998/test/"+Integer.toHexString(integer))));
            UnifiedBlock block=file.next();
            DataOutputStream dos=closer.register(new DataOutputStream(new BufferedOutputStream(new FileOutputStream(block.absolutePath()),8192)));
            startTime=System.nanoTime();
            for(int i=0;i<_1G/4;i++){
                dos.writeInt(rand.nextInt());
            }
            startTime=System.nanoTime()-startTime;
        }
        return startTime;
    }

    private static long test3(int integer)throws Exception{
        Random rand=new Random(integer);
        long startTime;
        try(Closer closer=Closer.create()){
            UnifiedFileSystem fs=UnifiedFileSystem.get(UnifiedFileSystem.newConfiguration(System.getenv("HADOOP_HOME")));
            UnifiedFile file=closer.register(fs.create(new URI("alluxio://hadoop06.ihep.ac.cn:19998/test/"+Integer.toHexString(integer))));
            UnifiedBlock block=file.next();
            MemoryBuffer mbuf=closer.register(new MemoryBuffer(new File(block.absolutePath()),0,_1G));
            startTime=System.nanoTime();
            for(int i=0;i<_1G/4;i++){
                mbuf.putInt(rand.nextInt());
            }
            startTime=System.nanoTime()-startTime;
        }
        return startTime;
    }

    public static void main(String args[])throws Exception{
        SparkConf conf=new SparkConf();
        try(JavaSparkContext jctx=new JavaSparkContext(conf); DataInputStream fis=new DataInputStream(new FileInputStream("/dev/urandom"))){
            int count=Integer.parseInt(args[0]);
            final int type=Integer.parseInt(args[1]);
            ArrayList<Integer> source=new ArrayList<>(count);
            for(int i=0;i<count;i++){
                source.add(fis.readInt());
            }
            JavaRDD<Integer> sourceRdd=jctx.parallelize(source,count);
            JavaRDD<Long> resultRdd=sourceRdd.map((Function<Integer, Long>) integer -> {
                switch (type){
                    case 0:
                        return test1(integer);
                    case 1:
                        return test2(integer);
                    default:
                        return test3(integer);
                }
            });
            List<Long> result=resultRdd.collect();
            double sumup=0;
            int sz=result.size();
            for (Long aResult : result) {
                sumup += aResult;
            }
            System.out.println("Total time: "+(sumup/sz/1e6)+"ms");
        }
    }
}
